import os
import shutil

from huggingface_hub import hf_hub_download
from modelscope import snapshot_download, model_file_download

preset_models_on_huggingface = {}
preset_models_on_modelscope = {}

website_preset_models = {
    "HuggingFace": preset_models_on_huggingface,
    "ModelScope": preset_models_on_modelscope,
}


class ModelFileDownloader(object):
    def __init__(self, priority, local_dir):
        self._priority = priority
        self._local_dir = local_dir
        os.makedirs(local_dir, exist_ok=True)
        self.download_fun = None
        self.select_download()

    def select_download(self):
        if self._priority == "HuggingFace":
            self.download_fun = hf_hub_download
        elif self._priority == "ModelScope":
            self.download_fun = snapshot_download

    def download(self, model_id, origin_file_path):
        file_name = os.path.basename(origin_file_path)
        if file_name in os.listdir(self._local_dir):
            print(f"{file_name} has been already in {self._local_dir}.")
        else:
            print(f"Start downloading {os.path.join(self._local_dir, file_name)}")
            if self._priority == "HuggingFace":
                hf_hub_download(model_id, origin_file_path, local_dir=self._local_dir)
                target_file_path = os.path.join(self._local_dir, file_name)
            elif self._priority == "ModelScope":
                snapshot_download(model_id, allow_file_pattern=origin_file_path, local_dir=self._local_dir)
                target_file_path = os.path.join(self._local_dir, os.path.split(origin_file_path)[-1])
            else:
                return
            downloaded_file_path = os.path.join(self._local_dir, origin_file_path)
            if downloaded_file_path != target_file_path:
                shutil.move(downloaded_file_path, target_file_path)
                shutil.rmtree(os.path.join(self._local_dir, origin_file_path.split("/")[0]))


class ModelLoader(object):
    def __init__(self, model_id_list: list, priority=["ModelScope", "HuggingFace"]):
        self._model_id_list = model_id_list
        self._priority = priority

    def load(self):
        downloaded_files = []
        load_files = []
        for model_id in self._model_id_list:
            for website in self._priority:
                if model_id in website_preset_models[website]:
                    # Parse model metadata
                    model_metadata = website_preset_models[website][model_id]
                    if isinstance(model_metadata, list):
                        file_data = model_metadata
                    else:
                        file_data = model_metadata.get("file_list", [])

                    # 从站点下载模型
                    model_files = []
                    for model_name, origin_file_path, local_dir in file_data:
                        # Check if the file is downloaded.
                        file_to_download = os.path.join(local_dir, os.path.basename(origin_file_path))
                        if file_to_download in downloaded_files:
                            continue
                        # Download
                        model_file_down = ModelFileDownloader(website, local_dir)
                        model_file_down.download(model_id, origin_file_path)
                        if os.path.basename(origin_file_path) in os.listdir(local_dir):
                            downloaded_files.append(file_to_download)
                            model_files.append(file_to_download)

                    # If the model is successfully downloaded, break.
                    if len(model_files) > 0:
                        if isinstance(model_metadata, dict) and "load_path" in model_metadata:
                            model_files = model_metadata["load_path"]
                        load_files.extend(model_files)
                        break

        return load_files
